"""
  验证跟随模型 transformer模型(生成验证数据)
"""


# !处理路径导入问题（添加绝对路径）！！！
import sys
import os
CODE_INTERNAL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')) # 生成Code文件夹内部对应的绝对路径
sys.path.append(CODE_INTERNAL_PATH)

# 导入外部包
import numpy as np
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import load_model

# 导入内部包
from utils.lstm import normalization, anit_normalization, create_lstm_dataset
from utils.read_data import read_extract_lanechange_data
from utils.write_data import write_data_to_file
from utils.transformer import PositionalEncoding

MODEL = "../model/lanechange_EPOCHS_50_BATCH_SIZE_64_transformer-normalization.keras"

FILE_PATH_I80_1_from = "../../../Data/Ngsim数据集/I80数据集/3. 提取数据/2. 换道数据/trajectories-0400-0415_lanechange.txt" # 300 80 21 / 过滤非法数据后 91 80 24（师兄原始的实验数据量 260 297 14）
FILE_PATH_I80_2_from = "../../../Data/Ngsim数据集/I80数据集/3. 提取数据/2. 换道数据/trajectories-0500-0515_lanechange.txt" # 251 80 21 / 过滤非法数据后 71 80 24（师兄原始的实验数据量 212 297 14）
FILE_PATH_I80_3_from = "../../../Data/Ngsim数据集/I80数据集/3. 提取数据/2. 换道数据/trajectories-0515-0530_lanechange.txt" # 281 80 21 / 过滤非法数据后 73 80 24（师兄原始的实验数据量 211 297 14）
FILE_PATH_101_1_from = "../../../Data/Ngsim数据集/101数据集/3. 提取数据/2. 换道数据/trajectories-0750am-0805am_lanechange.txt" # 293 80 21 / 过滤非法数据后 148 80 24（师兄原始的实验数据量 215 297 14）
FILE_PATH_101_2_from = "../../../Data/Ngsim数据集/101数据集/3. 提取数据/2. 换道数据/trajectories-0805am-0820am_lanechange.txt" # 300 80 21 / 过滤非法数据后 163 80 24（师兄原始的实验数据量 242 297 14）
FILE_PATH_101_3_from = "../../../Data/Ngsim数据集/101数据集/3. 提取数据/2. 换道数据/trajectories-0820am-0835am_lanechange.txt" # 326 80 21 / 过滤非法数据后 136 80 24（师兄原始的实验数据量 227 297 14）

FILE_PATH_to = "../../../Data/Ngsim数据集/I80数据集/4. 验证数据/2. 换道数据/trajectories-0400-0415_transformer_lanechange.txt" # 11458 5 10

def getData():
  data1 = read_extract_lanechange_data(FILE_PATH_I80_1_from)
  data2 = read_extract_lanechange_data(FILE_PATH_I80_2_from)
  data3 = read_extract_lanechange_data(FILE_PATH_I80_3_from)
  data4 = read_extract_lanechange_data(FILE_PATH_101_1_from)
  data5 = read_extract_lanechange_data(FILE_PATH_101_2_from)
  data6 = read_extract_lanechange_data(FILE_PATH_101_3_from)

  data = np.concatenate((data1, data2, data3, data4, data5, data6), axis=0)
  return data

def get_verify_data():
  # 读取数据
  datas = getData()

  # 归一化
  datas, min_max_list = normalization(datas, [10, 11, 12, 13, 14, 15, 18, 19, 21], [True, True, True, False, False, False, True, True, True])

  # 创建训练集和验证集
  n_past, n_future = 20, 5
  X, Y = create_lstm_dataset(datas, n_past, n_future, [10, 11, 12, 13, 14, 15, 18, 19, 21], [15]) # !历史是，预测是加速度
  x_train, x_val, y_train, y_val = train_test_split(X, Y, test_size=0.3, random_state=42)

  return x_val, y_val, min_max_list

def save_predict_data(y_val_true, y_val_pred, mean, std, lower_bound_1_std, upper_bound_1_std, lower_bound_2_std, upper_bound_2_std, lower_bound_3_std, upper_bound_3_std, target_fileName, min_max_list):
  print("valid_label: ", len(y_val_true), len(y_val_true[0]), len(y_val_true[0][0])) # 11458 5 10
  print("valid_predict: ", len(y_val_pred), len(y_val_pred[0]), len(y_val_pred[0][0])) # 11458 5 10
  print("mean", len(mean), len(mean[0]), len(mean[0][0])) # 11458 5 10
  print("std", len(std), len(std[0]), len(std[0][0])) # 11458 5 10
  print("lower_bound_1_std", len(lower_bound_1_std), len(lower_bound_1_std[0]), len(lower_bound_1_std[0][0]))
  
  # 反归一化
  y_val_true = anit_normalization(y_val_true, [0], [min_max_list[5]], [False])
  y_val_pred = anit_normalization(y_val_pred, [0], [min_max_list[5]], [False])

  # 先是真实数据，再是预测数据
  combine_data = np.concatenate((y_val_true, y_val_pred, mean, std, lower_bound_1_std, upper_bound_1_std, lower_bound_2_std, upper_bound_2_std, lower_bound_3_std, upper_bound_3_std), axis=2).tolist()
  write_data_to_file(combine_data, target_fileName)

  return combine_data

if __name__ == '__main__':
  # 获取训练模型时的数据
  x_val, y_val, min_max_list = get_verify_data()

  # 加载模型
  model = load_model(MODEL, custom_objects={'PositionalEncoding': PositionalEncoding})

  # 模型评估
  loss = model.evaluate(x_val, y_val, batch_size=64)
  y_val_predict = model(x_val)
  print("loss: ", loss) # 0.02091570571064949

  # 蒙特卡洛 Dropout 预测
  n_samples = 5
  mc_preds = np.stack([model(x_val, training=True) for _ in range(n_samples)])
  mc_preds = np.stack([anit_normalization(mc_preds[i], [0], [min_max_list[5]], [False]) for i in range(n_samples)]) # 反归一化
  mean = mc_preds.mean(axis=0) # 均值
  std = mc_preds.std(axis=0) # 方差（量化不确定性）
  print("mc_preds", len(mc_preds), len(mc_preds[0]), len(mc_preds[0][0]), len(mc_preds[0][0][0])) # 5 11458 5 10

  # 基于均值和方差计算下限
  lower_bound_1_std = mean - std
  upper_bound_1_std = mean + std
  lower_bound_2_std = mean - 2 * std
  upper_bound_2_std = mean + 2 * std
  lower_bound_3_std = mean - 3 * std
  upper_bound_3_std = mean + 3 * std

  # 保存预测数据
  # combine_data = save_predict_data(y_val, y_val_predict, mean, std, lower_bound_1_std, upper_bound_1_std, lower_bound_2_std, upper_bound_2_std, lower_bound_3_std, upper_bound_3_std, FILE_PATH_to, min_max_list)
